package pl.edu.pb.wi.pwnography.tree;

import java.io.Serializable;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import quickdt.Instance;

public class Tree {
    Set<Instance> instances;
    private Set<String> attributes;

    public Set<String> getAttributes() {
	return attributes;
    }

    public Tree(Set<Instance> instances) {
	this.instances = instances;

	for (Instance i : instances) {
	    if (attributes == null) {
		attributes = i.attributes.keySet();
		break;
	    }
	}
    }

    public Node build(Set<Instance> instances, Serializable attributeValue,
	    Set<String> attributes) {
	Node node = new Node();

	if (instances.isEmpty()) {
	    Map<Serializable, Integer> mostCommon = new HashMap<Serializable, Integer>();
	    int bestCount = 0;
	    Serializable bestClass = null;
	    for (Instance i : this.instances) {
		Integer val = null;
		if ((val = mostCommon.get(i.classification)) != null)
		    mostCommon.put(i.classification, ++val);
		else
		    mostCommon.put(i.classification, Integer.valueOf(1));

		if ((val = mostCommon.get(i.classification)) > bestCount) {
		    bestCount = val;
		    bestClass = i.classification;
		}
	    }

	    node.setResult(bestClass);
	    return node;
	}

	double classEntropy = Entropy.calculateClassEntropy(instances);
	if (classEntropy == 0) {
	    for (Instance i : instances) {
		node.setResult(i.classification);
		return node;
	    }
	}

	Map<Serializable, Integer> mostCommon = new HashMap<Serializable, Integer>();
	int bestCount = 0;
	Serializable bestClass = null;
	for (Instance i : instances) {
	    Integer val = null;
	    if ((val = mostCommon.get(i.classification)) != null)
		mostCommon.put(i.classification, ++val);
	    else
		mostCommon.put(i.classification, Integer.valueOf(1));

	    if ((val = mostCommon.get(i.classification)) > bestCount) {
		bestCount = val;
		bestClass = i.classification;
	    }
	}

	if (attributes.isEmpty() || mostCommon.size() == 1) {
	    node.setResult(bestClass);
	    return node;
	}

	String bestAttr = null;
	double bestEntropy = -1;
	for (String attr : attributes) {
	    double attrEntropy = classEntropy
		    - Entropy.calculateConditionalEntropy(instances, attr);
	    if (attrEntropy > bestEntropy) {
		bestEntropy = attrEntropy;
		bestAttr = attr;
	    }
	}

	node.setAttribute(bestAttr);
	node.setEntropy(bestEntropy);

	Map<Serializable, Set<Instance>> attrInstances = new HashMap<Serializable, Set<Instance>>();
	for (Instance i : instances) {
	    if (!attrInstances.containsKey(i.attributes.get(bestAttr)))
		attrInstances.put(i.attributes.get(bestAttr),
			new HashSet<Instance>());

	    attrInstances.get(i.attributes.get(bestAttr)).add(i);
	}
	for (Instance i : this.instances) {
	    if (!attrInstances.containsKey(i.attributes.get(bestAttr)))
		attrInstances.put(i.attributes.get(bestAttr),
			new HashSet<Instance>());
	}

	Set<String> nodeAttributes = new HashSet<String>(attributes);
	nodeAttributes.remove(bestAttr);
	for (Map.Entry<Serializable, Set<Instance>> attr : attrInstances
		.entrySet()) {
	    node.getChildren().put(attr.getKey(),
		    build(attr.getValue(), attr.getKey(), nodeAttributes));
	}

	return node;
    }
}
